# -*- coding:utf-8 -*-
"""
作者：XRF
日期：2023/10/10
IP = '127.0.0.1'时为本机IP代称
IP = '255.255.255.255'时为广播到本地局域网中
"""
import socket
import time

import cv2
import struct
import datetime
import math
import numpy as np


class VTDUdp:
    def __init__(self, ip, port, communicate_type=''):
        """
        Parameters
        ----------
        ip: ip地址
        port: 端口号
        communicate_type: 传输类型，接收图片或者接收信息
        """
        self.ADDR_img = (ip, port)  # 接收图片用端口: port
        self.ADDR_mes = (ip, port + 1)  # 接收信息用端口: port+1
        self.BUFFSIZE = 1024  # 接收xy_choose_ITL_D参数的缓冲区大小
        self.cv2img = None
        self.message = None
        self.uav_parameters = [0, 0, 0, 0]
        self.hasData = False
        self.finish = False  # 结束线程的标志

        # 以下两个均为不同用处的校验参数，注意必须为整数(没有更改的必要)
        self.CheckSum_img = 1234567890  # 图片传输时用
        self.CheckSum_mes = 1000010000  # 信息传输时用(未使用)

        # 专用于接收图片的socket
        self.socket_img = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)  # 创建UDP套接字
        self.socket_img.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 60000 * 100)  # 设定接收缓冲区大小

        # 普通socket
        self.socket_normal = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)  # 创建UDP套接字
        self.socket_normal.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)  # 允许发送广播数据
        self.socket_normal.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)  # 允许重用本地地址和端口

        # bind是服务器端(即接收信息端)调用的函数，绑定的是"自身"的IP
        # 假设服务端有多个网卡，即会有多个IP，这个时候我们可以选择用bind绑定其中一个IP，那么服务端只接收该IP上某端口的数据。
        # IP处填'0.0.0.0'表示绑定服务器端(即接收信息端)所有的IP地址，即表示所有发送到服务器的这个端口，不管是哪个网卡/哪个IP地址接收到的数据，都由这个服务端进程进行处理。
        if communicate_type == 'receive_img':
            self.socket_img.bind(('0.0.0.0', port))  # 接收图片用端口: port
        elif communicate_type == 'receive_mes':
            self.socket_normal.bind(('0.0.0.0', port + 1))  # 接收信息用端口: port+1
        elif communicate_type == '':
            print('普通发送信息')
        else:
            raise Exception('communicate_type参数有误，请从receive_img(接收图片)和receive_mes(接收信息)两者中选择')

    def udp_send_img(self, cv2img):
        """
        发出经过cv2处理后的图片
        Parameters
        ----------
        cv2img: 经过cv2处理后的图片
        """
        # 发送经过cv2读取后的图片到服务器
        img_encode = cv2.imencode(".png", cv2img)[1]
        data_encode = np.array(img_encode)
        data = data_encode.tostring()

        imgPackUnit = 60000
        imgLen = len(data)
        imgpackNum = imgLen // imgPackUnit + 1
        datetime_now = datetime.datetime.now()
        timeStmpSend = datetime.datetime.timestamp(datetime_now)

        # 循环发送图片码流
        for i in range(imgpackNum):
            if imgPackUnit * (i + 1) > len(data):  # 末尾包数据提取
                dataSend = data[imgPackUnit * i:]
            else:  # 前面数据包直接提取60000数据
                dataSend = data[imgPackUnit * i: imgPackUnit * (i + 1)]
            PackLen = 4 * 4 + 8 * 1 + len(dataSend)  # fhead的4i1d长度，加上图片数据长度
            fhead = struct.pack(
                "4i1d", self.CheckSum_img, PackLen, i, imgpackNum, timeStmpSend
            )  # 校验，包长度，包序号，总包数，时间戳
            dataSend = fhead + dataSend  # 包头加上图像数据
            self.socket_normal.sendto(dataSend, self.ADDR_img)  # 发送出去

    def udp_receive_img(self):
        """
        注意，本代码只适用于Config.json文件中TypeID=1，SendProtocol[0]=1的情况，且只能传输一张图
        接收经过cv2处理后的图片
        """
        fhead_size = struct.calcsize("4i1d")
        imgPackUnit = 60000
        seqList = []
        dataList = []

        recPackNum = 0
        timeStmpStore = 0
        while True:
            if self.finish:  # 结束线程的标志
                break

            try:
                buf, addr = self.socket_img.recvfrom(imgPackUnit + 2000)  # 加一些余量，确保包头数据考虑在内
            except socket.error:
                continue

            if len(buf) < fhead_size:
                print("数据包还没包头fhead_size长，数据错误")
                continue
            dd = struct.unpack("4i1d", buf[0:fhead_size])  # 校验，包长度，包序号，总包数，时间戳
            # print(dd)
            if dd[0] != self.CheckSum_img or dd[1] != len(buf):
                print("检验失败或数据长度不对")
                continue
            packSeq = dd[2]  # 包序号
            if packSeq == 0:  # 如果是第一个包
                seqList = []  # 清空数据序号列表
                dataList = []  # 清空数据缓存列表
                seqList = seqList + [packSeq]  # 提取序号
                dataList = dataList + [buf[fhead_size:]]  # 提取包头剩余数据
                timeStmpStore = dd[4]  # 提取时间戳
                recPackNum = dd[3]  # 以包头定义的总包作为接收结束标志
            else:  # 如果不是包头，直接将其存入列表
                if recPackNum == 0:
                    continue

                # 如果时间戳不一致
                if not math.isclose(timeStmpStore, dd[4], rel_tol=0.00001):
                    continue  # 跳过这个包
                seqList = seqList + [packSeq]  # 提取序号
                dataList = dataList + [buf[fhead_size:]]  # 提取包头剩余数据

            if len(seqList) == recPackNum:  # 如果收到的包达到总数了，开始处理图像
                recPackNum = 0
                data_total = b""
                dataOk = True
                for i in range(len(seqList)):
                    if seqList.count(i) < 1:
                        dataOk = False
                        print("处理img包失败，某序号不在包中")
                        break
                    idx0 = seqList.index(i)  # 按次序搜索包序号
                    data_total = data_total + dataList[idx0]

                if dataOk:  # 如果数据都没问题，开始处理图像
                    # 注意，本代码只适用于Config.json文件中TypeID=1，SendProtocol[0]=1的情况
                    nparr = np.frombuffer(data_total, np.uint8)
                    colorType = cv2.IMREAD_COLOR
                    self.cv2img = cv2.imdecode(nparr, colorType)
                    if self.cv2img is None:
                        print("图像不存在")
                        self.hasData = False
                    else:
                        self.hasData = True

    def udp_send_message(self, message):
        """
        发送xy_choose_ITL_D参数值
        :param
        message: xy_choose_ITL_D参数值
        :return:
        """
        if type(message) is list:
            message_body = ','.join(map(str, message))
        elif message == 'stop':
            message_body = message
        else:
            message_body = str(message)

        message_send = message_body
        self.socket_normal.sendto(message_send.encode('utf-8'), self.ADDR_mes)

    def udp_receive_message(self):
        """
        接收xy_choose_ITL_D参数值
        :return:
        self.message: xy_choose_ITL_D参数值
        """
        while True:
            if self.finish:  # 结束线程的标志
                break

            try:
                message_body = self.socket_normal.recvfrom(self.BUFFSIZE)[0].decode('utf-8')
            except socket.error:
                continue

            if message_body == 'None':
                self.message = None
            elif message_body == 'stop':
                self.message = message_body
                self.finish = True
            else:
                self.message = list(map(int, message_body.split(',')))

            print('已接收目标信息，' + str(self.message))

    def udp_send_uav_parameters(self, uav_parameters):
        """
        发送无人机的控制参数(向前、向右、向下、向右偏航的速度控制信号，即[vx, vy, vz, yawrate])
        :param
        uav_parameters: 无人机的控制参数 [vx, vy, vz, yawrate]
        :return:
        """
        if type(uav_parameters) is list and len(uav_parameters) == 4:
            parameters_body = ','.join(map(str, uav_parameters))
        else:
            print('无人机参数必须为列表形式，且长度必须为4')
            parameters_body = '0,0,0,0'

        parameters_send = parameters_body
        self.socket_normal.sendto(parameters_send.encode('utf-8'), self.ADDR_mes)

    def udp_receive_uav_parameters(self):
        """
        接收无人机的控制参数(向前、向右、向下、向右偏航的速度控制信号，即vx, vy, vz, yawrate)
        :return:
        self.uav_parameters: 无人机的控制参数 [vx, vy, vz, yawrate]
        """
        while True:
            if self.finish:  # 结束线程的标志
                break

            try:
                parameters_body = self.socket_normal.recvfrom(self.BUFFSIZE)[0].decode('utf-8')
            except socket.error:
                continue

            self.uav_parameters = list(map(float, parameters_body.split(',')))

    def close(self):
        """
        用于结束函数udp_receive_img的死循环
        """
        self.finish = True
        time.sleep(2)
        self.socket_img.close()
        self.socket_normal.close()


if __name__ == '__main__':
    # 是否为服务器
    server_type = True

    # 服务器执行代码
    if server_type:
        IP = 'localhost'  # 为发送信息的客户端IP
        PORT = 8888  # 服务器和客户端端口号必须相同，具体值可上网搜索参考


    # 客户端执行代码
    else:
        IP = 'localhost'  # 为接收信息的服务器IP
        PORT = 8888  # 服务器和客户端端口号必须相同，具体值可上网搜索参考



